/**
 * @file           : ekf.hpp
 * @author         : Star_Plucking
 * @brief          : 扩展卡尔曼滤波器
 * @date           : 25-5-3 下午7:45
 * @lastEditor     :
 * @lastEditTime   : 25-5-3 下午7:45
 */
#pragma once

// #include "std/function.h"
#include "functional"
#include <algorithm/matrix.h>

// 关于扩展卡尔曼和卡尔曼的区别如下：
// KF: x(k+1) = Ax(k) + Bu(k) + w(k), z(k) = Hx(k) + v(k)
// EKF: x(k+1) = f(x(k), u(k)) + w(k), z(k) = h(x(k)) + v(k)
// 这里我们认为卡尔曼是扩展卡尔曼的特例退化
namespace os::kalman {
// -----------------------扩展卡尔曼滤波器---------------------------
template <uint32_t state_dim, uint32_t measurement_dim = state_dim> // 状态维数和测量维数
class ExtensionKalmanFilter {
public:
    using StateVector = os::math::Matrix<state_dim, 1>; // 状态向量
    using StateMatrix = os::math::Matrix<state_dim, state_dim>; // 状态协方差矩阵
    using MeasurementVector = os::math::Matrix<measurement_dim, 1>; // 测量向量
    using MeasurementMatrix = os::math::Matrix<measurement_dim, measurement_dim>; // 测量协方差矩阵
    using KalmanGain = os::math::Matrix<state_dim, measurement_dim>; // 卡尔曼增益矩阵
    using JacobianF = os::math::Matrix<state_dim, state_dim>; // 状态转移雅可比矩阵(A阵)
    using JacobianH = os::math::Matrix<measurement_dim, state_dim>; // 测量雅可比矩阵(C阵列)

    // 状态转移函数类型：f(x, u, dt)
    using StateTransitionFn = std::function<StateVector(StateVector const&, float const&)>;

    // 测量函数类型：h(x)
    using MeasurementFn = std::function<MeasurementVector(StateVector const&)>;

    // 雅可比矩阵计算函数类型 这两个雅各比正是扩展卡尔曼线性化的关键
    using StateJacobianFn = std::function<JacobianF(StateVector const&, float const&)>;
    using MeasurementJacobianFn = std::function<JacobianH(StateVector const&)>;

    /**
     * @brief 构造函数
     * @param f 状态转移函数
     * @param h 测量函数
     * @param jF 状态转移雅可比矩阵计算函数
     * @param jH 测量雅可比矩阵计算函数
     */
    ExtensionKalmanFilter(StateTransitionFn f, MeasurementFn h, StateJacobianFn jF, MeasurementJacobianFn jH)
        : f_(f), h_(h), jacobian_f_(jF), jacobian_h_(jH) {
        // 初始化为单位矩阵乘以一个小系数
        P_ = math::getIdentityMatrix<state_dim>();
        Q_ = math::getIdentityMatrix<state_dim>() * 0.01f;
        R_ = math::getIdentityMatrix<measurement_dim>() * 0.01f;

        // 初始化状态向量为零
        float data[measurement_dim]{};
        x_ = StateVector(data);
    }

    /**
     * @brief 初始化滤波器状态
     * @param x0 初始状态
     * @param P0 初始协方差
     */
    void initialize(StateVector const& x0, StateMatrix const& P0) {
        x_ = x0;
        P_ = P0;
    }

    /**
     * @brief 设置过程噪声协方差
     * @param Q 过程噪声协方差矩阵
     */
    void setProcessNoise(StateMatrix const& Q) { Q_ = Q; }

    /**
     * @brief 设置测量噪声协方差
     * @param R 测量噪声协方差矩阵
     */
    void setMeasurementNoise(MeasurementMatrix const& R) { R_ = R; }

    /**
     * @brief 预测步骤
     * @param dt 时间增量
     */
    void predict(float const& dt) {
        // 非线性状态预测
        x_ = f_(x_, dt);

        // 计算状态转移雅可比矩阵
        JacobianF F = jacobian_f_(x_, dt);

        // 更新协方差
        P_ = F * P_ * F.T() + Q_;
    }

    /**
     * @brief 更新步骤
     * @param z 测量值
     */
    void update(MeasurementVector const& z) {
        // 计算测量预测
        MeasurementVector z_pred = h_(x_);

        // 计算测量雅可比矩阵
        JacobianH H = jacobian_h_(x_);

        // 计算测量残差
        MeasurementVector y = z - z_pred;

        // 计算残差协方差
        MeasurementMatrix S = H * P_ * H.T() + R_;

        // 计算卡尔曼增益
        KalmanGain K = P_ * H.T() * S.inverse();

        // 更新状态
        x_ = x_ + K * y;

        // 更新协方差 (使用Joseph形式以保证正定性)
        // 在实际应用中，特别是在长期运行的系统或数值精度受限的嵌入式系统中，
        // Joseph形式能够显著提高滤波器的稳定性和可靠性，尽管它的计算量略大于标准形式。（红豆泥？）
        StateMatrix I = os::math::getIdentityMatrix<state_dim>();
        P_ = (I - K * H) * P_ * (I - K * H).T() + K * R_ * K.T();
    }

    /**
     * @brief 获取当前状态
     * @return 当前状态向量
     */
    StateVector const& getState() const { return x_; }

    /**
     * @brief 获取当前状态协方差
     * @return 当前状态协方差矩阵
     */
    StateMatrix const& getCovariance() const { return P_; }

private:
    StateVector x_; // 状态向量
    StateMatrix P_; // 状态协方差
    StateMatrix Q_; // 过程噪声协方差
    MeasurementMatrix R_; // 测量噪声协方差

    StateTransitionFn f_; // 状态转移函数
    MeasurementFn h_; // 测量函数
    StateJacobianFn jacobian_f_; // 状态转移雅可比矩阵计算函数
    MeasurementJacobianFn jacobian_h_; // 测量雅可比矩阵计算函数
};

// -------------------标准卡尔曼---------------------
template <uint32_t state_dim, uint32_t measurement_dim = state_dim, uint32_t control_dim = 0>
class KalmanFilter {
public:
    using StateVector = os::math::Matrix<state_dim, 1>;
    using StateMatrix = os::math::Matrix<state_dim, state_dim>;
    using MeasurementVector = os::math::Matrix<measurement_dim, 1>;
    using MeasurementMatrix = os::math::Matrix<measurement_dim, measurement_dim>;
    using StateTransitionMatrix = os::math::Matrix<state_dim, state_dim>; // A矩阵
    using MeasurementTransitionMatrix = os::math::Matrix<measurement_dim, state_dim>; // H矩阵
    using ControlMatrix = os::math::Matrix<state_dim, control_dim>; // B矩阵

    /**
     * @brief 构造函数，使用现有的扩展卡尔曼生成
     * @param A 状态转移矩阵
     * @param H 测量矩阵
     * @param B 控制输入矩阵（可选）
     */
    KalmanFilter(StateTransitionMatrix const& A, MeasurementTransitionMatrix const& H,
                 ControlMatrix const& B = ControlMatrix().zeros())
        : A_(A),
          H_(H),
          B_(B),
          ekf_(
              // 状态转移函数 f(x, dt) = A*x + B*u (这里简化为没有控制输入)
              [this](StateVector const& x, float const& dt) -> StateVector { return A_ * x; },
              // 测量函数 h(x) = H*x
              [this](StateVector const& x) -> MeasurementVector { return H_ * x; },
              // 状态转移雅可比矩阵 (线性系统中就是A矩阵本身)
              [this](StateVector const& x, float const& dt) -> StateTransitionMatrix { return A_; },
              // 测量雅可比矩阵 (线性系统中就是H矩阵本身)
              [this](StateVector const& x) -> MeasurementTransitionMatrix { return H_; }) {
    }

    // 直接使用ExtensionKalmanFilter提供的方法
    void initialize(StateVector const& x0, StateMatrix const& P0) { ekf_.initialize(x0, P0); }

    void setProcessNoise(StateMatrix const& Q) { ekf_.setProcessNoise(Q); }

    void setMeasurementNoise(MeasurementMatrix const& R) { ekf_.setMeasurementNoise(R); }

    void predict(float const& dt) { ekf_.predict(dt); }

    void update(MeasurementVector const& z) { ekf_.update(z); }

    StateVector const& getState() const { return ekf_.getState(); }

    StateMatrix const& getCovariance() const { return ekf_.getCovariance(); }

private:
    StateTransitionMatrix A_; // 状态转移矩阵
    MeasurementTransitionMatrix H_; // 测量矩阵
    ControlMatrix B_; // 控制输入矩阵
    ExtensionKalmanFilter<state_dim, measurement_dim> ekf_; // 内部使用扩展卡尔曼滤波器
};
} // namespace os::kalman